from utils.log import Logger
import pandas as pd
# from imblearn.over_sampling import SMOTE
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, roc_auc_score, precision_score, recall_score, accuracy_score
import joblib
import os
import warnings
from sklearn.exceptions import ConvergenceWarning

# 抑制LogisticRegression的收敛警告
warnings.filterwarnings('ignore', category=ConvergenceWarning)
warnings.filterwarnings('ignore', message='lbfgs failed to converge')

# 创建日志记录器
logger = Logger(root_path=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 
                log_name='train', level='info').get_logger()


def model_setup(data: pd.DataFrame, target_col: str):
    """
    建立模型
    输入：1.数据集 2.标签列名称
    将输入的数据集进行模型训练
    并将保存训练好的模型
    """
    try:
        logger.info(f"开始模型训练，数据形状: {data.shape}，目标列: {target_col}")
        
        # 分离特征和标签
        y = data[target_col]
        x = data.drop(columns=target_col)
        logger.info(f"特征数量: {x.shape[1]}，样本数量: {x.shape[0]}")
        logger.info(f"目标变量分布: {y.value_counts().to_dict()}")
        
        # 创建并训练模型
        estimator = LogisticRegression(C=20.0)
        x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=13)
        estimator.fit(x_train, y_train)
        logger.info("模型训练完成")
        
        # 对测试集进行预测和评估
        y_test_pred = estimator.predict(x)
        y_test_proba = estimator.predict_proba(x)[:, 1]
        test_auc = roc_auc_score(y, y_test_proba)
        
        # 记录测试集AUC值
        logger.info(f"测试集AUC值: {test_auc:.4f}")
        
        # 记录测试集分类报告
        test_report = classification_report(y, y_test_pred)
        logger.info(f"测试集分类报告:\n{test_report}")
        
        # 保存模型
        model_dir = '../models'
        os.makedirs(model_dir, exist_ok=True)
        model_path = os.path.join(model_dir, 'NB_model.pkl')
        joblib.dump(estimator, model_path)
        logger.info(f"模型已保存到: {model_path}")
        
        logger.info("模型训练流程完成")
        return
        
    except KeyError as e:
        logger.error(f"数据中缺少目标列: {e}")
        raise
    except ValueError as e:
        logger.error(f"数据格式错误: {e}")
        raise
    except Exception as e:
        logger.error(f"模型训练过程中出现未知错误: {e}")
        raise
